package com.flink.hbase.kafka2elasticsearch.timepartition;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.datatype.jsr310.JavaTimeModule;
import org.apache.flink.streaming.api.functions.ProcessFunction;
import org.apache.flink.util.Collector;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.Serializable;
import java.time.Instant;
import java.time.LocalDateTime;
import java.time.ZoneId;
import java.time.format.DateTimeFormatter;
import java.time.format.DateTimeParseException;
import java.util.regex.Pattern;

/**
 * 业务时间提取器
 * 从 Kafka 消息中提取业务时间，用于时间分区
 */
public class BusinessTimeExtractor implements Serializable {
    private static final Logger LOG = LoggerFactory.getLogger(BusinessTimeExtractor.class);
    private static final long serialVersionUID = 1L;
    
    private final String[] timeFields;
    private final String defaultTimeZone;
    private final boolean useEventTime;
    private final boolean useProcessingTime;
    
    private transient ObjectMapper objectMapper;
    private transient ZoneId zoneId;
    
    // 支持的时间格式
    private static final DateTimeFormatter[] DATE_FORMATTERS = {
        DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss"),
        DateTimeFormatter.ofPattern("yyyy-MM-dd'T'HH:mm:ss"),
        DateTimeFormatter.ofPattern("yyyy-MM-dd'T'HH:mm:ss.SSS"),
        DateTimeFormatter.ofPattern("yyyy-MM-dd'T'HH:mm:ss.SSSX"),
        DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss.SSS"),
        DateTimeFormatter.ISO_LOCAL_DATE_TIME,
        DateTimeFormatter.ISO_OFFSET_DATE_TIME,
        DateTimeFormatter.ISO_ZONED_DATE_TIME
    };
    
    // 时间戳格式检测
    private static final Pattern TIMESTAMP_PATTERN = Pattern.compile("^\\d{10,13}$");
    
    public BusinessTimeExtractor(String[] timeFields, String defaultTimeZone, 
                               boolean useEventTime, boolean useProcessingTime) {
        this.timeFields = timeFields != null ? timeFields : new String[]{"timestamp", "create_time", "update_time", "event_time"};
        this.defaultTimeZone = defaultTimeZone != null ? defaultTimeZone : "Asia/Shanghai";
        this.useEventTime = useEventTime;
        this.useProcessingTime = useProcessingTime;
    }
    
    /**
     * 从消息中提取业务时间
     */
    public LocalDateTime extractBusinessTime(String message) {
        if (objectMapper == null) {
            initializeMapper();
        }
        
        // 如果配置使用处理时间，直接返回当前时间
        if (useProcessingTime) {
            return LocalDateTime.now(zoneId);
        }
        
        try {
            // 尝试解析 JSON 消息
            JsonNode jsonNode = objectMapper.readTree(message);
            
            // 按优先级顺序尝试提取时间字段
            for (String timeField : timeFields) {
                LocalDateTime businessTime = extractTimeFromField(jsonNode, timeField);
                if (businessTime != null) {
                    LOG.debug("Extracted business time {} from field {}", businessTime, timeField);
                    return businessTime;
                }
            }
            
            // 如果配置使用事件时间但没有找到，使用当前时间
            if (useEventTime) {
                LOG.warn("No business time found in message, using current time");
                return LocalDateTime.now(zoneId);
            }
            
        } catch (Exception e) {
            LOG.error("Error extracting business time from message: {}", message, e);
        }
        
        // 默认使用当前时间
        return LocalDateTime.now(zoneId);
    }
    
    /**
     * 从 JSON 字段中提取时间
     */
    private LocalDateTime extractTimeFromField(JsonNode jsonNode, String fieldName) {
        JsonNode fieldNode = getNestedField(jsonNode, fieldName);
        if (fieldNode == null || fieldNode.isNull()) {
            return null;
        }
        
        String timeValue = fieldNode.asText();
        if (timeValue == null || timeValue.trim().isEmpty()) {
            return null;
        }
        
        return parseTimeValue(timeValue.trim());
    }
    
    /**
     * 获取嵌套字段（支持 a.b.c 格式）
     */
    private JsonNode getNestedField(JsonNode jsonNode, String fieldPath) {
        String[] paths = fieldPath.split("\\.");
        JsonNode currentNode = jsonNode;
        
        for (String path : paths) {
            if (currentNode == null || !currentNode.has(path)) {
                return null;
            }
            currentNode = currentNode.get(path);
        }
        
        return currentNode;
    }
    
    /**
     * 解析时间值
     */
    private LocalDateTime parseTimeValue(String timeValue) {
        // 尝试解析时间戳
        if (TIMESTAMP_PATTERN.matcher(timeValue).matches()) {
            return parseTimestamp(timeValue);
        }
        
        // 尝试解析各种日期时间格式
        for (DateTimeFormatter formatter : DATE_FORMATTERS) {
            try {
                return LocalDateTime.parse(timeValue, formatter);
            } catch (DateTimeParseException e) {
                // 继续尝试下一个格式
            }
        }
        
        LOG.warn("Unable to parse time value: {}", timeValue);
        return null;
    }
    
    /**
     * 解析时间戳
     */
    private LocalDateTime parseTimestamp(String timestampStr) {
        try {
            long timestamp = Long.parseLong(timestampStr);
            
            // 判断是秒级还是毫秒级时间戳
            if (timestamp > 1000000000000L) {
                // 毫秒级时间戳
                return LocalDateTime.ofInstant(Instant.ofEpochMilli(timestamp), zoneId);
            } else {
                // 秒级时间戳
                return LocalDateTime.ofInstant(Instant.ofEpochSecond(timestamp), zoneId);
            }
        } catch (NumberFormatException e) {
            LOG.warn("Invalid timestamp format: {}", timestampStr);
            return null;
        }
    }
    
    /**
     * 初始化 ObjectMapper 和时区
     */
    private void initializeMapper() {
        objectMapper = new ObjectMapper();
        objectMapper.registerModule(new JavaTimeModule());
        zoneId = ZoneId.of(defaultTimeZone);
    }
    
    /**
     * ProcessFunction 实现，用于在数据流中提取业务时间
     */
    public static class TimeExtractionFunction extends ProcessFunction<String, TimePartitionedMessage> {
        private static final long serialVersionUID = 1L;
        
        private final BusinessTimeExtractor timeExtractor;
        private final String baseIndexName;
        
        public TimeExtractionFunction(BusinessTimeExtractor timeExtractor, String baseIndexName) {
            this.timeExtractor = timeExtractor;
            this.baseIndexName = baseIndexName;
        }
        
        @Override
        public void processElement(String message, Context ctx, Collector<TimePartitionedMessage> out) throws Exception {
            try {
                LocalDateTime businessTime = timeExtractor.extractBusinessTime(message);
                TimePartitionInfo partitionInfo = TimePartitionInfo.of(baseIndexName, businessTime);
                
                TimePartitionedMessage partitionedMessage = new TimePartitionedMessage(
                    message, businessTime, partitionInfo
                );
                
                out.collect(partitionedMessage);
                
            } catch (Exception e) {
                LOG.error("Error processing message for time extraction: {}", message, e);
                // 使用当前时间作为默认值
                LocalDateTime currentTime = LocalDateTime.now();
                TimePartitionInfo partitionInfo = TimePartitionInfo.of(baseIndexName, currentTime);
                
                TimePartitionedMessage partitionedMessage = new TimePartitionedMessage(
                    message, currentTime, partitionInfo
                );
                
                out.collect(partitionedMessage);
            }
        }
    }
    
    /**
     * 带时间分区信息的消息
     */
    public static class TimePartitionedMessage implements Serializable {
        private static final long serialVersionUID = 1L;
        
        private final String originalMessage;
        private final LocalDateTime businessTime;
        private final TimePartitionInfo partitionInfo;
        
        public TimePartitionedMessage(String originalMessage, LocalDateTime businessTime, TimePartitionInfo partitionInfo) {
            this.originalMessage = originalMessage;
            this.businessTime = businessTime;
            this.partitionInfo = partitionInfo;
        }
        
        public String getOriginalMessage() {
            return originalMessage;
        }
        
        public LocalDateTime getBusinessTime() {
            return businessTime;
        }
        
        public TimePartitionInfo getPartitionInfo() {
            return partitionInfo;
        }
        
        @Override
        public String toString() {
            return "TimePartitionedMessage{" +
                    "businessTime=" + businessTime +
                    ", partitionInfo=" + partitionInfo +
                    ", messageLength=" + (originalMessage != null ? originalMessage.length() : 0) +
                    '}';
        }
    }
    
    /**
     * 构建器
     */
    public static class Builder {
        private String[] timeFields;
        private String defaultTimeZone = "Asia/Shanghai";
        private boolean useEventTime = true;
        private boolean useProcessingTime = false;
        
        public Builder timeFields(String... timeFields) {
            this.timeFields = timeFields;
            return this;
        }
        
        public Builder defaultTimeZone(String defaultTimeZone) {
            this.defaultTimeZone = defaultTimeZone;
            return this;
        }
        
        public Builder useEventTime(boolean useEventTime) {
            this.useEventTime = useEventTime;
            return this;
        }
        
        public Builder useProcessingTime(boolean useProcessingTime) {
            this.useProcessingTime = useProcessingTime;
            return this;
        }
        
        public BusinessTimeExtractor build() {
            return new BusinessTimeExtractor(timeFields, defaultTimeZone, useEventTime, useProcessingTime);
        }
    }
    
    /**
     * 创建构建器
     */
    public static Builder builder() {
        return new Builder();
    }
    
    /**
     * 创建默认配置的时间提取器
     */
    public static BusinessTimeExtractor createDefault() {
        return new Builder().build();
    }
} 